#!/bin/bash
set -eo pipefail

CONFIG_FILE="$HOME/.docker_mirrors"
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 初始化配置文件
init_config() {
    mkdir -p "$(dirname "$CONFIG_FILE")"
    [ -f "$CONFIG_FILE" ] || echo "# Docker镜像加速站列表（每行一个）" > "$CONFIG_FILE"
}

# 显示带颜色的错误信息
error_exit() {
    echo -e "${RED}错误: $1${NC}" >&2
    exit 1
}

# 处理set命令
handle_set() {
    local accelerator="$1"
    
    # 去重并添加到文件头部
    grep -vFx "$accelerator" "$CONFIG_FILE" > tmp_config
    mv tmp_config "$CONFIG_FILE"
    sed -i "1i$accelerator" "$CONFIG_FILE"
    echo -e "${GREEN}✅ 已添加加速站: ${YELLOW}$accelerator${NC}（已置顶）${NC}"
}

# 智能路径补全函数
smart_path() {
    local image="$1"
    
    # 自动添加 library/ 前缀
    if [[ "$image" != */* ]]; then
        image="library/$image"
    fi

    echo "$image"
}

# 处理pull命令
handle_pull() {
    local image_tag="$1"
    
    # 获取加速器列表
    accelerators=$(cat "$CONFIG_FILE")
    [ -z "$accelerators" ] && error_exit "未配置加速站，请先使用 'set' 命令添加"

    failed=()
    success=0


    # 智能补全路径
    full_image_name=$(smart_path "$image_tag")

    # 尝试每个加速器
    for accelerator in $accelerators; do
        echo -e "\n${YELLOW}尝试加速站: ${accelerator}${NC}"
        
        # 构建完整镜像地址
        target_image="${accelerator}/${full_image_name}"


        echo -e "\n${YELLOW}拉取命令: ${target_image}${NC}"
        if stdbuf -oL docker pull "$target_image" 2>&1 | tee /dev/tty; then
            echo -e "${GREEN}✔ 成功拉取！正在重命名...${NC}"
            docker tag "$target_image" "$image_tag"
            docker rmi "$target_image"
            success=1
            break
        else
            failed+=("$accelerator")
            echo -e "${RED}! 拉取失败: $accelerator${NC}"
        fi
    done

    # 最终结果展示
    echo -e "\n${GREEN}=== 操作结果 ===${NC}"
    if [ $success -eq 1 ]; then
        echo -e "${GREEN}✓ 成功获取镜像: ${image_tag}${NC}"
    else
        echo -e "${RED}✗ 所有加速站均失败，尝试的地址: ${failed[*]}${NC}"
        error_exit "所有加速站均不可用"
    fi
}

# 清空加速站列表
handle_clear() {
    if [ -f "$CONFIG_FILE" ]; then
        rm -f "$CONFIG_FILE"
        echo -e "${GREEN}✅ 已清空所有加速站配置${NC}"
    else
        echo -e "${YELLOW}ℹ 配置文件不存在${NC}"
    fi
}

main() {
    trap 'rm -f tmp_config' EXIT

    case "$1" in
        set)
            [[ -z "$2" ]] && error_exit "用法: $0 set <镜像加速站地址>"
            handle_set "$2"
            ;;
        pull)
            [[ -z "$2" ]] && error_exit "用法: $0 pull <镜像名称>"
            handle_pull "$2"
            ;;
        clear)
            handle_clear
            ;;
        *)
            echo -e "${YELLOW}用法:"
            echo "  $0 set <加速站地址>       添加镜像加速站（无需协议）"
            echo "  $0 pull <镜像名称>        直接拉取任意格式镜像"
            echo "  $0 clear                 清空所有加速站配置"
            echo -e "${NC}"
            exit 1
            ;;
    esac
}

# 初始化运行
init_config
main "$@"